//给你一棵无根带权树，树中总共有 n 个节点，分别表示 n 个服务器，服务器从 0 到 n - 1 编号。同时给你一个数组 edges ，其中 edges[
//i] = [ai, bi, weighti] 表示节点 ai 和 bi 之间有一条双向边，边的权值为 weighti 。再给你一个整数 signalSpeed 。
// 
//
// 如果两个服务器 a ，b 和 c 满足以下条件，那么我们称服务器 a 和 b 是通过服务器 c 可连接的 ： 
//
// 
// a < b ，a != c 且 b != c 。 
// 从 c 到 a 的距离是可以被 signalSpeed 整除的。 
// 从 c 到 b 的距离是可以被 signalSpeed 整除的。 
// 从 c 到 b 的路径与从 c 到 a 的路径没有任何公共边。 
// 
//
// 请你返回一个长度为 n 的整数数组 count ，其中 count[i] 表示通过服务器 i 可连接 的服务器对的 数目 。 
//
// 
//
// 示例 1： 
//
// 
//
// 
//输入：edges = [[0,1,1],[1,2,5],[2,3,13],[3,4,9],[4,5,2]], signalSpeed = 1
//输出：[0,4,6,6,4,0]
//解释：由于 signalSpeed 等于 1 ，count[c] 等于所有从 c 开始且没有公共边的路径对数目。
//在输入图中，count[c] 等于服务器 c 左边服务器数目乘以右边服务器数目。
// 
//
// 示例 2： 
//
// 
//
// 
//输入：edges = [[0,6,3],[6,5,3],[0,3,1],[3,2,7],[3,1,6],[3,4,2]], signalSpeed = 3
//输出：[2,0,0,0,0,0,2]
//解释：通过服务器 0 ，有 2 个可连接服务器对(4, 5) 和 (4, 6) 。
//通过服务器 6 ，有 2 个可连接服务器对 (4, 5) 和 (0, 5) 。
//所有服务器对都必须通过服务器 0 或 6 才可连接，所以其他服务器对应的可连接服务器对数目都为 0 。
// 
//
// 
//
// 提示： 
//
// 
// 2 <= n <= 1000 
// edges.length == n - 1 
// edges[i].length == 3 
// 0 <= ai, bi < n 
// edges[i] = [ai, bi, weighti]
// 
// 1 <= weighti <= 10⁶ 
// 1 <= signalSpeed <= 10⁶ 
// 输入保证 edges 构成一棵合法的树。 
// 
//
// Related Topics 树 深度优先搜索 数组 👍 17 👎 0


package LeetCode.editor.cn;

import java.util.*;
import java.util.stream.IntStream;

/**
 * @author ldltd
 * @date 2024-06-04 11:17:31
 * @description 3067.在带权树网络中统计可连接服务器对数目
 */
public class CountPairsOfConnectableServersInAWeightedTreeNetwork{
	 public static void main(String[] args) {
	 	 //测试代码
	 	 CountPairsOfConnectableServersInAWeightedTreeNetwork fun=new CountPairsOfConnectableServersInAWeightedTreeNetwork();
	 	 Solution solution = fun.new Solution();
		  //solution.countPairsOfConnectableServers(new int[][]{{0,1,1},{1,2,5},{2,3,13},{3,4,9},{4,5,2}},1);
//		 solution.countPairsOfConnectableServers(new int[][]{{0,6,3},{6,5,3},{0,3,1},{3,2,7},{3,1,6},{3,4,2}},3);
		 solution.countPairsOfConnectableServers(new int[][]{{1,0,2},{2,1,4},{3,2,4},{4,0,3},{5,1,4},{6,2,2},{7,6,4},{8,1,2},{9,8,3}},1);
	 }
	 
//力扣代码
//leetcode submit region begin(Prohibit modification and deletion)
class Solution {
    public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {
		Set<Integer> vis=new HashSet<>();
		List<List<Integer []>> ed=new ArrayList<>();
		// 初始化邻接表
		for (int[] edge : edges) {
			while (ed.size() <= Math.max(edge[0], edge[1])) {
				ed.add(new ArrayList<>());
			}
			ed.get(edge[0]).add(new Integer[]{edge[1], edge[2]});
			ed.get(edge[1]).add(new Integer[]{edge[0], edge[2]});
		}

		/*  for (int[] edge : edges) {
            graph.computeIfAbsent(edge[0], k -> new ArrayList<>()).add(new int[]{edge[1], edge[2]});
            graph.computeIfAbsent(edge[1], k -> new ArrayList<>()).add(new int[]{edge[0], edge[2]});
        }*/
		int[] res = new int[ed.size()];
		for (int i=0;i<res.length;i++) {
			if(ed.get(i).size()<2){
				res[i]=0;
				vis.remove(i);
			}else {
				List<Integer> t=new ArrayList<>();
				for (Integer []j : ed.get(i)) {
					vis.clear();
					vis.add(j[0]);
					vis.add(i);
					t.add(dfs(ed,j[0],j[1],vis, signalSpeed)+(j[1]%signalSpeed==0?1:0));
				}
				res[i]=helper(t);
			}
		}
		return res;
    }
	public int dfs(List<List<Integer[]>> ed,int i,int count,Set<Integer> memo,int signalSpeed){
			int p=0;
			for (Integer[] idx : ed.get(i)) {
				if(!memo.contains(idx[0])){
					memo.add(idx[0]);
					if((count + idx[1])%signalSpeed==0) p++;
					int dfs = dfs(ed, idx[0], count + idx[1], memo, signalSpeed);
					p+=dfs;
				}
			}
		return p;
	}
	public int helper(List<Integer> t){
		if(t.size()<2) return 0;

//		return IntStream.range(0, t.size())
//				.flatMap(i -> IntStream.range(0, t.size())
//						.filter(j -> i != j)
//						.map(j -> t.get(i) * t.get(j)))
//				.sum() / 2; // 因为每对 (i, j) 被计算了两次，所以除以 2
		return IntStream.range(0, t.size())
				.flatMap(i -> IntStream.range(i + 1, t.size())
						.map(j -> t.get(i) * t.get(j)))
				.sum();
	}



	public int[] countPairsOfConnectableServers1(int[][] edges, int signalSpeed) {
		int n = edges.length + 1;
		List<int[]>[] g = new ArrayList[n];
		Arrays.setAll(g, i -> new ArrayList<>());
		//保存图
		for (int[] e : edges) {
			int x = e[0];
			int y = e[1];
			int wt = e[2];
			g[x].add(new int[]{y, wt});
			g[y].add(new int[]{x, wt});
		}

		int[] ans = new int[n];
		for (int i = 0; i < n; i++) {
			//比如每个节点为中间服务器，对应的每条边有3,4,6个符合条件
			//则3*4与4*6与3*6两两都可以
			//可以简化成 3*4-》3+4-》7*6，用sum来表示这个过程
			int sum = 0;
			for (int[] e : g[i]) {
				int cnt = dfs(e[0], i, e[1], g, signalSpeed);
				ans[i] += cnt * sum;
				sum += cnt;
			}
		}
		return ans;
	}
	// x为中心点临街的点,fa为中心点，sum为累计的边权,g是图，signalSpeed是除数
	private int dfs(int x, int fa, int sum, List<int[]>[] g, int signalSpeed) {
		int cnt = sum % signalSpeed == 0 ? 1 : 0;
		for (int[] e : g[x]) {
			int y = e[0];
			if (y != fa) {
				cnt += dfs(y, x, sum + e[1], g, signalSpeed);
			}
		}
		return cnt;
	}

}
//leetcode submit region end(Prohibit modification and deletion)

}
